import numpy as np   
import pandas as pd 
import os

class Get_feature_data():
    
    def __init__(self, args) -> None:
        self.args = args
        if args.feature_root is None:
            self.args.feature_root = 'data/cmodty/my_feature/feature_data'
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
    
    def get_feature(self):  
        feature_list = []
        for i in self.args.feature_list:
            path = os.path.join(self.args.feature_root, i)
            feature_i = pd.read_pickle(path)
            if self.args.varieties is not None:
                feature_i = feature_i[self.args.varieties]
            feature_i = feature_i.fillna(method='pad', axis=0)
            feature_i = feature_i[(feature_i.index >= self.beg_time) & (feature_i.index <= self.end_time)]
            feature_list.append((feature_i.values).astype(float))
            print(f'获取特征 {i} 完成')
        codes = feature_i.columns
        times = np.array(feature_i.index)
        feature = np.stack(feature_list,axis=2)
        print('合并特征完成')
        return feature, codes, times    

class Get_npz_feature():
    def __init__(self, args) -> None:
        self.args = args
        if self.args.feature_root is None:
            self.args.feature_root = 'data/cmodty/orgdata/comdty.npz'
        self.data = np.load(self.args.feature_root)
        self.datetime_list = pd.to_datetime([f"{str(date)}0{str(time)}" for date in self.data['dates'] for time in self.data['times']])
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
      
    def get_feature(self):
        feature_list = []
        codes = self.data['codes']
        for i in self.args.feature_list:
            feature_i = self.df_make(self.data[i])
            if self.args.varieties is not None:
                feature_i = feature_i[self.args.varieties]
            feature_i = feature_i[(feature_i.index >= self.beg_time) & (feature_i.index <= self.end_time)]
            times = np.array(feature_i.index)
            feature_list.append(feature_i.values)
            print(f'获取特征 {i} 完成')
        times = np.array(feature_i.index)
        feature = np.array(feature_list) #time*codes*feature_num
        feature = feature.swapaxes(0,2)
        feature = feature.swapaxes(0,1)
        print('合并特征完成')
        return feature, codes, times
    
    def df_make(self, data):
        org_shape = data.shape
        data = data.reshape(-1, org_shape[2])
        data_df = pd.DataFrame(data)
        data_df.index = pd.to_datetime(self.datetime_list)
        data_df.columns = self.data['codes']
        return data_df


class Get_feature_data_many():
    
    def __init__(self, args) -> None:
        self.args = args
        if args.feature_root is None:
            self.args.feature_root = ['data/cmodty/my_feature/feature_data']
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
    
    def get_feature(self): 
        name_lst = [] 
        feature_list = []
        for j in self.args.feature_root:
            lst = sorted(os.listdir(j))
            for i in lst:
                path = os.path.join(j, i)
                feature_i = pd.read_pickle(path).astype(np.float32)
                if self.args.varieties is not None:
                    feature_i = feature_i[self.args.varieties]
                feature_i = feature_i.replace([np.inf, -np.inf], np.nan)
                # feature_i = feature_i.fillna(method='pad', axis=0)
                feature_i = feature_i[(feature_i.index >= self.beg_time) & (feature_i.index <= self.end_time)]
                fv = feature_i.values
                
                fv[np.isinf(fv)] = np.nan
                nan_ratio = np.isnan(fv).sum()/ (fv.shape[0]*fv.shape[1])
                if (nan_ratio>0.8) or (np.nanstd(fv)==0): #(nan_ratio>0.8) or 
                    print('drop feature', i[:-9])
                    continue
                            
                feature_list.append((fv).astype(float))
                name_lst.append(i[:-9])
                print(f'获取特征 {i} 完成',feature_i.shape, np.isnan(feature_i.values.astype(float)).sum())
        codes = feature_i.columns
        times = np.array(feature_i.index)
        feature = np.stack(feature_list,axis=2)
        print('合并特征完成')
        return feature, codes, times , name_lst 

import json
class Get_feature_data_select():
    select_file = r'data\cmodty\check\feature_fliter_result.json'
    
    def __init__(self, args) -> None:
        self.args = args
        if args.feature_root is None:
            self.args.feature_root = ['data/cmodty/my_feature/feature_data']
        self.beg_time = self.args.time_param['insample_beg']
        self.end_time = self.args.time_param['outsample_end']
        
        with open(Get_feature_data_select.select_file,'r', encoding='UTF-8') as f:
            self.select_dict = json.load(f)
                    
    def get_feature(self): 
        select_lst = self.select_dict['reamain']
        name_lst = [] 
        feature_list = []
        for j in self.args.feature_root:
            lst = sorted(os.listdir(j))
            for i in lst:
                if i[:-9] in select_lst:
                    path = os.path.join(j, i)
                    feature_i = pd.read_pickle(path).astype(float)
                    if self.args.varieties is not None:
                        feature_i = feature_i[self.args.varieties]
                    # feature_i = feature_i.fillna(method='pad', axis=0)
                    feature_i = feature_i[(feature_i.index >= self.beg_time) & (feature_i.index <= self.end_time)]
                    fv = feature_i.values
                    feature_list.append((fv).astype(float))
                    name_lst.append(i[:-9])
                    print(f'获取特征 {i} 完成',feature_i.shape, np.isnan(feature_i.values.astype(float)).sum())
        codes = feature_i.columns
        times = np.array(feature_i.index)
        feature = np.stack(feature_list,axis=2)
        print('合并特征完成')
        return feature, codes, times , name_lst

    def get_tail(self):
        tail_lst = []
        for i in self.feature_root:
            tail_i = os.path.split(i)[-1]
            tail_lst.append((tail_i.split('_')[-1]))
        tail = ''
        for i in tail_lst:
            tail = tail+i
        return tail